Skip to content

[externalization] Add a two-step patch/sub-export API for module externalization - #53

Open
gokulkrishna98 wants to merge 3 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/module-externalization-api
Open

[externalization] Add a two-step patch/sub-export API for module externalization#53
gokulkrishna98 wants to merge 3 commits into
apple:mainfrom
gokulkrishna98:dev/gokul/module-externalization-api

Conversation

@gokulkrishna98

Copy link
Copy Markdown
Contributor

Reworks module externalization — emitting an nn.Module subtree as a single Core AI composite op — into an explicit two-step API, and adds autograd support so gradients can flow through externalized submodules.

The one-shot path is unchanged: add_pytorch_module(model, export_fn=…, externalize_modules=[...]) still does everything. The new split exists for callers who must drive export themselves (e.g. a quantizer), so marking and sub-export can happen either side of their own passes:

_patch_model_for_externalization(model, [spec])
ep = quantizer.prepare(model).calibrate(data).finalize()   # any caller-driven export
externalized = _subexport_and_restore(model, ep)

program = (
    TorchConverter()
    .add_exported_program(ep, _externalized_exported_programs=externalized)
    .to_coreai()
)

Changes

  • Pipeline documented as four phases: Mark → Prepare → Export → Emit. Restore is folded
    into _subexport_and_restore's finally, so the model is always left unpatched.
  • register_autograd: backward through an externalized submodule now works.

Port the source-side changes of the module externalization API into
coreai_torch/, replacing the marker-based approach.

Tests cover:
- backward: gradients flow through externalized submodules via
  register_autograd.
- no call sites: exporting before patching leaves no custom-op call
  sites, so _subexport_and_restore warns and skips externalization
  instead of raising, and conversion still succeeds.
- the manual patch/sub-export/convert workflow, i.e. the documented
  advanced path where the caller drives export themselves instead of
  using add_pytorch_module. A passthrough export stands in for the
  quantizer, pinning the workflow contract: custom-op call sites
  survive the caller's export, the model is left unpatched after
  _subexport_and_restore, and the marked submodule is emitted as a
  rms_norm composite op.
- numerics for that workflow: compiles the converted program, runs it
  in the Core AI runtime, and compares against the restored PyTorch
  model.
Cross-reference the two "no call sites" warning paths so neither reads as
dead code left over from the refactor: _drop_missing_call_sites handles
top-level marked submodules missing from the whole-model program, while
_PreparedModules.__iter__ handles the nested case it cannot see, once the
enclosing module's sub-export exists. Both docstrings now point at each
other, and _subexport_and_restore names the split.

Add test_externalize_multiple_staged_entries_numerics: two staged entries
in one to_coreai() call, one via add_pytorch_module(externalize_modules=)
and one plain add_exported_program, validating numerics on both
entrypoints. This pins the invariant that externalization is per-entry —
_init_conversion_state() resets _externalized_exported_programs before
each entry, so a plain entry converted alongside an externalized one
stays flat. Verified non-vacuous: dropping that reset and adding a third
_perform_externalization call site fails the test on a duplicate
rms_norm symbol.

_validate_numerics takes a function_name argument so it can check a
named entrypoint instead of only "main".
The module docstring documents four Phases (Mark, Prepare, Export, Emit)
while individual docstrings still used a separate six-Step scheme that no
longer lined up — export was "Step 4" but Phase 3, prepare was "Step 3"
but Phase 2, and there was no Step 2 at all. Relabel every per-function
reference to the Phase vocabulary and drop the number from
_restore_externalized, since restore is no longer its own phase.

Docstrings only; no behavior change.
@gokulkrishna98
gokulkrishna98 marked this pull request as ready for review August 4, 2026 16:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants